package org.budo.elasticsearch.jdbc.driver.statement;

import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.budo.elasticsearch.jdbc.driver.BudoElasticsearchJdbcConnection;
import org.budo.elasticsearch.jdbc.driver.BudoElasticsearchJdbcPreparedStatement;
import org.budo.elasticsearch.jdbc.driver.parser.jsql.JSqlParserStatement;
import org.elasticsearch.action.update.UpdateRequestBuilder;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.search.SearchHit;

import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.statement.update.Update;

/**
 * @author limingwei
 */
public class UpdateStatement extends AbstractElasticsearchStatement {
    @Override
    public void execute(PreparedStatement preparedStatement) throws SQLException {
        SearchHit[] searchHits = this.executeSearchRequest(preparedStatement, new String[] { "_id" });

        Map<String, Object> updateSetDoc = this.updateSetDoc(preparedStatement);

        for (SearchHit searchHit : searchHits) {
            this.updateDocById(preparedStatement, searchHit.getId(), updateSetDoc);
        }

        ((BudoElasticsearchJdbcPreparedStatement) preparedStatement).setUpdateCount(searchHits.length);
    }

    private Map<String, Object> updateSetDoc(PreparedStatement preparedStatement) {
        Update update = (Update) ((JSqlParserStatement) this.getSqlStatement()).getStatement();
        Map<String, Object> doc = new HashMap<String, Object>();
        List<Column> columns = update.getColumns();
        List<Expression> expressions = update.getExpressions();

        for (int i = 0; i < columns.size(); i++) {
            String columnName = columns.get(i).getColumnName();
            Expression expression = expressions.get(i);
            Object value = this.expressionValue(preparedStatement, expression);
            doc.put(columnName, value);
        }
        return doc;
    }

    private void updateDocById(PreparedStatement preparedStatement, String id, Map<String, Object> doc) throws SQLException {
        BudoElasticsearchJdbcConnection connection = (BudoElasticsearchJdbcConnection) preparedStatement.getConnection();
        TransportClient transportClient = connection.getTransportClient();

        UpdateRequestBuilder updateRequestBuilder = transportClient.prepareUpdate() //
                .setIndex(connection.getCatalog()) //
                .setType(this.tableName()) //
                .setId(id) //
                .setDoc(doc);

        this.executeAndGet(updateRequestBuilder);
    }
}